import xarray as xr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import QuantileTransformer
from sklearn.compose import TransformedTargetRegressor

from sklearn.neural_network import MLPRegressor
from sklearn.ensemble import AdaBoostRegressor
from sklearn.ensemble import BaggingRegressor

from sklearn.metrics import root_mean_squared_error as rmse

from tqdm import tqdm

import dill
import random

import salishsea_tools.viz_tools as sa_vi

def datasets_preparation(dataset, dataset2):
    
    drivers = np.stack([np.ravel(dataset['Temperature_(0m-15m)']),
        np.ravel(dataset['Temperature_(15m-100m)']), 
        np.ravel(dataset['Salinity_(0m-15m)']),
        np.ravel(dataset['Salinity_(15m-100m)']),
        np.ravel(dataset2['Summation_of_solar_radiation']),
        np.ravel(dataset2['Mean_wind_speed']),
        np.ravel(dataset2['Mean_air_temperature'])
        ])
    indx = np.where(~np.isnan(drivers).any(axis=0))
    drivers = drivers[:,indx[0]]

    diat = np.ravel(dataset['Diatom'])
    diat = diat[indx[0]]

    return(drivers, diat, indx)

def regressor (inputs, targets):
    
    inputs = inputs.transpose()

    model = TransformedTargetRegressor(regressor=make_pipeline(StandardScaler(),MLPRegressor(alpha=0.001, learning_rate='invscaling', epsilon=1e-07, hidden_layer_sizes=[100,100])), 
        transformer=QuantileTransformer())
    regr = BaggingRegressor(model, n_estimators=10, n_jobs=10).fit(inputs,targets)
    
    return(regr)

def regressor2 (inputs, targets, variable_name):

    inputs2 = inputs.transpose()

    outputs_test = regr.predict(inputs2)
   
    m = scatter_plot(targets, outputs_test, variable_name) 
    r = np.round(np.corrcoef(targets, outputs_test)[0][1],3)
    rms = rmse(targets, outputs_test)

    return (r, rms, m)

def regressor3 (inputs, targets):
  
    inputs2 = inputs.transpose()

    outputs_test = regr.predict(inputs2)
   
    # compute slope m and intercept b
    m, b = np.polyfit(targets, outputs_test, deg=1)
    
    r = np.round(np.corrcoef(targets, outputs_test)[0][1],3)
    rms = rmse(targets, outputs_test)

    return (r, rms, m)

def regressor4 (inputs, targets, variable_name):

    inputs2 = inputs.transpose()
    
    outputs = regr.predict(inputs2)

    # Post processing
    indx2 = np.full((len(diat_i.y)*len(diat_i.x)),np.nan)
    indx2[indx[0]] = outputs
    model = np.reshape(indx2,(len(diat_i.y),len(diat_i.x)))

    m = scatter_plot(targets, outputs, variable_name + str(dates[i].date())) 

    # Preparation of the dataarray 
    model = xr.DataArray(model,
        coords = {'y': diat_i.y, 'x': diat_i.x},
        dims = ['y','x'],
        attrs=dict( long_name = variable_name + "Concentration",
        units="mmol m-2"),)
                        
    plotting3(targets, model, diat_i, variable_name)

def printing (targets, outputs, m):

    print ('The amount of data points is', outputs.size)
    print ('The slope of the best fitting line is ', np.round(m,3))
    print ('The correlation coefficient is:', np.round(np.corrcoef(targets, outputs)[0][1],3))
    print (' The root mean square error is:', rmse(targets,outputs))

def scatter_plot(targets, outputs, variable_name):

    # compute slope m and intercept b
    m, b = np.polyfit(targets, outputs, deg=1)

    printing(targets, outputs, m)

    fig, ax = plt.subplots(2, figsize=(5,10), layout='constrained')

    ax[0].scatter(targets,outputs, alpha = 0.2, s = 10)

    lims = [np.min([ax[0].get_xlim(), ax[0].get_ylim()]),
        np.max([ax[0].get_xlim(), ax[0].get_ylim()])]

    # plot fitted y = m*x + b
    ax[0].axline(xy1=(0, b), slope=m, color='r')

    ax[0].set_xlabel('targets')
    ax[0].set_ylabel('outputs')
    ax[0].set_xlim(lims)
    ax[0].set_ylim(lims)
    ax[0].set_aspect('equal')

    ax[0].plot(lims, lims,linestyle = '--',color = 'k')

    h = ax[1].hist2d(targets,outputs, bins=100, cmap='jet', 
        range=[lims,lims], cmin=0.1, norm='log')
    
    ax[1].plot(lims, lims,linestyle = '--',color = 'k')

    # plot fitted y = m*x + b
    ax[1].axline(xy1=(0, b), slope=m, color='r')

    ax[1].set_xlabel('targets')
    ax[1].set_ylabel('outputs')
    ax[1].set_aspect('equal')

    fig.colorbar(h[3],ax=ax[1], location='bottom')

    fig.suptitle(variable_name)

    plt.show()

    return (m)

def plotting(variable, name):

    plt.plot(years,variable, marker = '.', linestyle = '')
    plt.xlabel('Years')
    plt.ylabel(name)
    plt.show()

def plotting2(variable,title):
    
    fig, ax = plt.subplots()

    scatter= ax.scatter(dates,variable, marker='.', c=pd.DatetimeIndex(dates).month)

    ax.legend(handles=scatter.legend_elements()[0], labels=['February','March','April'])
    fig.suptitle('Daily ' + title + ' (15 Feb - 30 Apr)')
    
    fig.show()

def plotting3(targets, model, variable, variable_name):

    fig, ax = plt.subplots(2,2, figsize = (10,15))

    cmap = plt.get_cmap('cubehelix')
    cmap.set_bad('gray')

    variable.plot(ax=ax[0,0], cmap=cmap, vmin = targets.min(), vmax =targets.max(), cbar_kwargs={'label': variable_name + ' Concentration  [mmol m-2]'})
    model.plot(ax=ax[0,1], cmap=cmap, vmin = targets.min(), vmax = targets.max(), cbar_kwargs={'label': variable_name + ' Concentration  [mmol m-2]'})
    ((variable-model) / variable * 100).plot(ax=ax[1,0], cmap=cmap, cbar_kwargs={'label': variable_name + ' Concentration  [percentage]'})

    plt.subplots_adjust(left=0.1,
        bottom=0.1, 
        right=0.95, 
        top=0.95, 
        wspace=0.35, 
        hspace=0.35)

    sa_vi.set_aspect(ax[0,0])
    sa_vi.set_aspect(ax[0,1])
    sa_vi.set_aspect(ax[1,0])


    ax[0,0].title.set_text(variable_name + ' (targets)')
    ax[0,1].title.set_text(variable_name + ' (outputs)')
    ax[1,0].title.set_text('targets - outputs')
    ax[1,1].axis('off')

    fig.suptitle(str(dates[i].date()))

    plt.show()

ds = xr.open_dataset('/data/ibougoudis/MOAD/files/integrated_model_var_old.nc')
ds2 = xr.open_dataset('/data/ibougoudis/MOAD/files/external_inputs.nc')

# ds = ds.isel(time_counter = (np.arange(0, len(ds.time_counter),2)), 
#     y=(np.arange(ds.y[0], ds.y[-1], 5)), 
#     x=(np.arange(ds.x[0], ds.x[-1], 5)))

# ds2 = ds2.isel(time_counter = (np.arange(0, len(ds2.time_counter),2)), 
#     y=(np.arange(ds2.y[0], ds2.y[-1], 5)), 
#     x=(np.arange(ds2.x[0], ds2.x[-1], 5)))

dataset = ds.sel(time_counter = slice('2007', '2020'))
dataset2 = ds2.sel(time_counter = slice('2007', '2020'))

drivers, diat, _ = datasets_preparation(dataset, dataset2)

regr = regressor(drivers, diat)

ds = ds.sel(time_counter = slice('2021', '2023'))
ds2 = ds2.sel(time_counter = slice('2021', '2023'))

dates = pd.DatetimeIndex(ds['time_counter'].values)

maps = random.sample(range(0,len(ds.time_counter)),10)

for i in tqdm(maps):

    dataset = ds.isel(time_counter=i)
    dataset2 = ds2.isel(time_counter=i)
    drivers, diat, indx = datasets_preparation(dataset, dataset2)

    diat_i = dataset['Diatom']

    regressor4(drivers, diat, 'Diatom ')

  0%|          | 0/10 [00:00<?, ?it/s]

The amount of data points is 46479
The slope of the best fitting line is  0.279
The correlation coefficient is: 0.43
 The root mean square error is: 0.20322512036762408

 10%|█         | 1/10 [00:15<02:22, 15.84s/it]

The amount of data points is 46479
The slope of the best fitting line is  0.525
The correlation coefficient is: 0.693
 The root mean square error is: 0.05155939498026469

 20%|██        | 2/10 [00:29<01:57, 14.71s/it]

The amount of data points is 46479
The slope of the best fitting line is  0.079
The correlation coefficient is: 0.051
 The root mean square error is: 0.2110631584483521

 30%|███       | 3/10 [00:42<01:37, 13.88s/it]

The amount of data points is 46479
The slope of the best fitting line is  0.477
The correlation coefficient is: 0.773
 The root mean square error is: 0.06722725285034359

 40%|████      | 4/10 [00:54<01:17, 12.98s/it]

The amount of data points is 46479
The slope of the best fitting line is  0.191
The correlation coefficient is: 0.218
 The root mean square error is: 0.15083785204669917

years = range (2021,2024)

for year in tqdm(range (2021,2024)):
    
    dataset = ds.sel(time_counter=str(year))
    dataset2 = ds2.sel(time_counter=str(year))

    drivers, diat, _ = datasets_preparation(dataset, dataset2)

    r, rms, m = regressor2(drivers, diat, 'Diatom ' + str(year))

  0%|          | 0/3 [00:00<?, ?it/s]/tmp/ipykernel_14146/1980467486.py:4: RankWarning: Polyfit may be poorly conditioned
  m, b = np.polyfit(targets, outputs, deg=1)

The amount of data points is 3485925
The slope of the best fitting line is  0.447
The correlation coefficient is: 0.645
 The root mean square error is: 0.1378248039117271

 33%|███▎      | 1/3 [01:19<02:39, 79.64s/it]/tmp/ipykernel_14146/1980467486.py:4: RankWarning: Polyfit may be poorly conditioned
  m, b = np.polyfit(targets, outputs, deg=1)

The amount of data points is 3485925
The slope of the best fitting line is  0.47
The correlation coefficient is: 0.557
 The root mean square error is: 0.12422049887596191

 67%|██████▋   | 2/3 [02:27<01:12, 72.79s/it]/tmp/ipykernel_14146/1980467486.py:4: RankWarning: Polyfit may be poorly conditioned
  m, b = np.polyfit(targets, outputs, deg=1)

The amount of data points is 3485925
The slope of the best fitting line is  0.5
The correlation coefficient is: 0.371
 The root mean square error is: 0.16468230161096423

100%|██████████| 3/3 [03:36<00:00, 72.16s/it]
100%|██████████| 3/3 [03:36<00:00, 72.16s/it]

r_all2 = np.array([])
rms_all2 = np.array([])
slope_all2 = np.array([])

for i in tqdm(range (0, len(ds.time_counter))):
    
    dataset = ds.isel(time_counter=i)
    dataset2 = ds2.isel(time_counter=i)

    drivers, diat, _ = datasets_preparation(dataset, dataset2)

    r, rms, m = regressor3(drivers, diat)

    r_all2 = np.append(r_all2,r)
    rms_all2 = np.append(rms_all2,rms)
    slope_all2 = np.append(slope_all2,m)

plotting2(r_all2, 'Correlation Coefficients')
plotting2(rms_all2, 'Root Mean Square Errors')
plotting2(slope_all2, 'Slope of the best fitting line')

100%|██████████| 225/225 [35:02<00:00,  9.34s/it]

Script to reproduce years based on a model trained with random points¶

Importing¶

Datasets Preparation¶

Regressor¶

Regressor 2¶

Regressor 3¶

Regressor 4¶

Printing¶

Scatter Plot¶

Plotting¶

Plotting 2¶

Plotting 3¶

Training (Random Points)¶

Daily Maps¶

Other Years (Anually)¶

Other Years (Daily)¶